############################################################################
# 2019 - present Contributed by Apulis Technology (Shenzhen) Co. LTD
#
# This program and the accompanying materials are made available under the
# terms of the MIT License, which is available at
# https://www.opensource.org/licenses/MIT
#
# See the NOTICE file distributed with this work for additional
# information regarding copyright ownership.
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
#
# SPDX-License-Identifier: MIT
############################################################################

execpath=$(dirname $(realpath $0))
export RANK_ID=0
export DEVICE_ID=0

export ASCEND_GLOBAL_LOG_LEVEL=0
export TF_CPP_MIN_LOG_LEVEL=0
export SLOG_PRINT_TO_STDOUT=0

if [[ -z $VISIBLE_IDS ]]; then
  export VISIBLE_IDS=$ASCEND_VISIBLE_DEVICES
fi

OLD_IFS="$IFS"
IFS=","
DEVICE_IDS=($VISIBLE_IDS)
IFS=$OLD_IFS

echo "DEVICE_IDS---${DEVICE_IDS[@]}"

export RANK_INDEX=0
export HCCL_CONNECT_TIMEOUT=600
export RANK_START=0
export DEVICE_NUM=${#DEVICE_IDS[@]}
export RANK_SIZE=${#DEVICE_IDS[@]}
unset POD_NAME

ulimit -u unlimited

#if [ -f MNIST.zip ]; then
#  unzip MNIST.zip
#  rm -rf *.zip
#fi

unzip MNIST.zip

rm -rf *ckpt* *meta *.zip
if [[ -z "$VISIBLE_IDS" ]]; then
  # 如果不在平台环境中，直接使用0,1卡，使用默认的table_file和环境变量
  DEVICE_IDS=(0 1)
  export RANK_TABLE_FILE=$execpath/hccl_2p.json
  export JOB_ID=0
  export DEVICE_NUM=${#DEVICE_IDS[@]}
  export RANK_SIZE=${#DEVICE_IDS[@]}
#  source $execpath/../env_21.0.2.sh

# 如果在平台环境中
else
  # 如果在是分布式环境
  if [[ $TASK_NUM_WORKER -gt 1 ]]; then
    export SERVER_ID=${TASK_ROLE_IDX}
    export RANK_START=$((DEVICE_NUM * SERVER_ID))
    export RANK_SIZE=$((TASK_NUM_WORKER * 8))
  fi
fi

echo "========================"
echo "start tensorflow testcase"
echo "========================"

rm -rf /tmp/$JOB_ID
#tf与ms现在已经更新为相同启动方式
for ((device_phy_id = 0; device_phy_id < ${DEVICE_NUM}; device_phy_id++)); do
  export RANK_ID=$((RANK_START + device_phy_id))
  export RANK_INDEX=$((RANK_START + device_phy_id))
  export DEVICE_ID=$device_phy_id
  export DEVICE_INDEX=$device_phy_id
  export ASCEND_DEVICE_ID=$device_phy_id

  mkdir -p /tmp/$JOB_ID/device_$device_phy_id
  cd /tmp/$JOB_ID/device_$device_phy_id

  echo DEVICE_ID-$DEVICE_ID
  echo RANK_ID-$RANK_ID
  python $execpath/tf-lenet.py $execpath &
done
wait

echo "========================"
echo "finish tensorflow testcase"
echo "========================"

echo "========================"
echo "start mindspore testcase"
echo "========================"

rm -rf /tmp/$JOB_ID
for ((device_phy_id = 0; device_phy_id < ${DEVICE_NUM}; device_phy_id++)); do
  export RANK_ID=$((RANK_START + device_phy_id))
  export RANK_INDEX=$((RANK_START + device_phy_id))
  export DEVICE_ID=$device_phy_id
  export DEVICE_INDEX=$device_phy_id
  export ASCEND_DEVICE_ID=$device_phy_id

  mkdir -p /tmp/$JOB_ID/device_$device_phy_id
  cd /tmp/$JOB_ID/device_$device_phy_id

  echo DEVICE_ID-$DEVICE_ID
  echo RANK_ID-$RANK_ID
  python $execpath/ms-lenet.py --data_path=$execpath/MNIST &
done
wait

echo "========================"
echo "start torch testcase"
echo "========================"

#python $execpath/torch-lenet.py $execpath

echo "========================"
echo "finish pytorch testcase"
echo "========================"

if [[ $TASK_NUM_WORKER -gt 1 ]]; then
  if [[ "$VCJOB_TASK_NAME" != "master" ]]; then
    echo "npu worker completed , wait for master !"
    sleep infinity
  fi
fi
